package org.simplesocket.simplenettycore.utils;

import org.reflections.Reflections;
import org.reflections.scanners.Scanners;
import org.reflections.util.ClasspathHelper;
import org.reflections.util.ConfigurationBuilder;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.OAEPParameterSpec;
import javax.crypto.spec.PSource;
import javax.crypto.spec.SecretKeySpec;
import java.io.File;
import java.io.IOException;
import java.lang.annotation.Annotation;
import java.lang.reflect.Constructor;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.JarURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.MGF1ParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

public class SimpleNettyUtils {
    
    public static class ClassUtil {
        
        public static List<Class<?>> scan(String packageName) {
            // 1. 获取包的路径，并转换为文件路径格式
            String path = packageName.replace('.', '/');
            ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
            if (classLoader == null) {
                return null;
            }
            
            // 2. 找到资源
            Enumeration<URL> resources;
            try {
                resources = classLoader.getResources(path);
            } catch (IOException e) {
                return null;
            }
            List<Class<?>> classes = new ArrayList<>();
            
            // 3. 遍历所有找到的资源
            while (resources.hasMoreElements()) {
                URL resource = resources.nextElement();
                String protocol = resource.getProtocol();
                
                if ("file".equals(protocol)) {
                    try {
                        // 文件系统中的类处理
                        classes.addAll(findClassesInFileSystem(new File(resource.getFile()), packageName));
                    }catch (ClassNotFoundException ignore){
                        return null;
                    }
                } else if ("jar".equals(protocol)) {
                    try {
                        // JAR文件中的类处理
                        JarURLConnection jarConnection = (JarURLConnection) resource.openConnection();
                        JarFile jarFile = jarConnection.getJarFile();
                        classes.addAll(findClassesInJarFile(jarFile, path));
                    }catch (Exception ignore){
                        return null;
                    }
                }
            }
            return classes;
        }
        
        public static <T extends Annotation> Set<Class<?>> scanAnnotatedClasses(String packageName, Class<T> annotation){
            Reflections ref = new Reflections(packageName);
            return ref.getTypesAnnotatedWith(annotation);
        }
        
        public static <T extends Annotation> Set<Method> scanAnnotatedMethods(String packageName, Class<T> annotation){
            ConfigurationBuilder config = new ConfigurationBuilder()
                    .setUrls(ClasspathHelper.forPackage(packageName))
                    .setScanners(Scanners.MethodsAnnotated, Scanners.SubTypes);
            Reflections ref = new Reflections(config);
            return ref.getMethodsAnnotatedWith(annotation);
        }
        
        public static <T extends Annotation> Set<Field> scanAnnotatedFields(String packageName, Class<T> annotation){
            ConfigurationBuilder config = new ConfigurationBuilder()
                    .setUrls(ClasspathHelper.forPackage(packageName))
                    .setScanners(Scanners.FieldsAnnotated);
            Reflections ref = new Reflections(config);
            return ref.getFieldsAnnotatedWith(annotation);
        }
        
        public static String getUserBasePackageName() {
            StackTraceElement[] stackTrace = Thread.currentThread().getStackTrace();
            try {
                for (StackTraceElement element : stackTrace) {
                    String className = element.getClassName();
                    if (!className.startsWith("org.simplesocket.simplenettycore") && !className.startsWith("java.lang.Thread")){
                        Class<?> callerClass = Class.forName(className);
                        return callerClass.getPackage().getName();
                    }
                }
            }catch (Exception e){
                throw new RuntimeException("Can not find base packages,because:\n" + e);
            }
            
            throw new RuntimeException("Can not find base packages ...");
        }
        
        // 文件系统中获取类
        private static List<Class<?>> findClassesInFileSystem(File directory, String packageName) throws ClassNotFoundException {
            List<Class<?>> classes = new ArrayList<>();
            if (!directory.exists()) {
                return classes;
            }
            File[] files = directory.listFiles();
            if (files != null) {
                for (File file : files) {
                    if (file.isDirectory()) {
                        // 递归处理子目录
                        classes.addAll(findClassesInFileSystem(file, packageName + "." + file.getName()));
                    } else if (file.getName().endsWith(".class")) {
                        // 这里的".class"文件去掉后缀
                        String className = packageName + '.' + file.getName().substring(0, file.getName().length() - 6);
                        classes.add(Class.forName(className));
                    }
                }
            }
            return classes;
        }
        
        // 从JAR文件中获取类
        private static List<Class<?>> findClassesInJarFile(JarFile jarFile, String path) throws ClassNotFoundException {
            List<Class<?>> classes = new ArrayList<>();
            Enumeration<JarEntry> entries = jarFile.entries(); // 获取JAR文件中的所有条目
            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement();
                String entryName = entry.getName();
                if (entryName.startsWith(path) && entryName.endsWith(".class") && !entryName.contains("$")) {
                    // 类的完整路径并替换文件分隔符
                    String className = entryName.substring(0, entryName.length() - 6).replace('/', '.');
                    classes.add(Class.forName(className));
                }
            }
            return classes;
        }
        
        public static <T> T getSingleton(Class<T> clazz) throws NoSuchMethodException, InvocationTargetException,
                InstantiationException, IllegalAccessException {
            Constructor<T> constructor = clazz.getDeclaredConstructor();
            constructor.setAccessible(true);
            return clazz.cast(constructor.newInstance());
        }
        
    }
    
    public static class AESUtil {
        
        private static final String ALGORITHM = "AES/CBC/PKCS5Padding";
        
        // 生成指定位数的AES密钥
        public static SecretKey generateKey(int keySize) throws NoSuchAlgorithmException {
            KeyGenerator keyGen = KeyGenerator.getInstance("AES");
            keyGen.init(keySize);
            return keyGen.generateKey();
        }
        
        // 生成指定位数的AES密钥，并返回Base64编码的字符串
        public static String generateKeyAsBase64(int keySize) throws NoSuchAlgorithmException {
            SecretKey key = generateKey(keySize);
            return Base64.getEncoder().encodeToString(key.getEncoded());
        }
        
        // 从字符串生成AES密钥
        public static SecretKey generateKeyFromString(String keyString) {
            byte[] decodedKey = Base64.getDecoder().decode(keyString);
            return new SecretKeySpec(decodedKey, 0, decodedKey.length, "AES");
        }
        
        // 生成随机的32字符密钥字符串（适用于256位AES）
        public static String generateRandomKeyString() {
            SecureRandom random = new SecureRandom();
            byte[] keyBytes = new byte[32];
            random.nextBytes(keyBytes);
            return Base64.getEncoder().encodeToString(keyBytes);
        }

        public static String encrypt(String plaintext, String base64Key) throws Exception {
            // 将Base64密钥解码为字节数组
            byte[] keyBytes = Base64.getDecoder().decode(base64Key);
            if (keyBytes.length != 32) {
                throw new IllegalArgumentException("Key must be 32 bytes long (256 bits)");
            }

            // 生成随机的IV（初始化向量）
            byte[] ivBytes = new byte[16];
            new SecureRandom().nextBytes(ivBytes);
            IvParameterSpec iv = new IvParameterSpec(ivBytes);

            // 初始化AES加密
            SecretKeySpec secretKey = new SecretKeySpec(keyBytes, "AES");
            Cipher cipher = Cipher.getInstance(ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, secretKey, iv);

            // 加密数据
            byte[] encryptedBytes = cipher.doFinal(plaintext.getBytes(StandardCharsets.UTF_8));

            // 返回 IV:EncryptedData 格式的字符串
            String ivBase64 = Base64.getEncoder().encodeToString(ivBytes);
            String encryptedBase64 = Base64.getEncoder().encodeToString(encryptedBytes);
            return ivBase64 + ":" + encryptedBase64;
        }
        
        public static String decrypt(String encryptedText, String base64Key) throws Exception {
            // 将Base64密钥解码为字节数组
            byte[] keyBytes = Base64.getDecoder().decode(base64Key);
            if (keyBytes.length != 32) {
                throw new IllegalArgumentException("Key must be 32 bytes long (256 bits)");
            }
            
            // 解码密文，并分离出IV和实际的加密内容
            String[] parts = encryptedText.split(":");
            byte[] ivBytes = Base64.getDecoder().decode(parts[0]);
            byte[] encryptedBytes = Base64.getDecoder().decode(parts[1]);
            
            // 初始化AES解密
            SecretKeySpec secretKey = new SecretKeySpec(keyBytes, "AES");
            IvParameterSpec iv = new IvParameterSpec(ivBytes);
            
            Cipher cipher = Cipher.getInstance(ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, secretKey, iv);
            
            // 解密并返回明文
            byte[] decryptedBytes = cipher.doFinal(encryptedBytes);
            return new String(decryptedBytes, StandardCharsets.UTF_8);
        }
        
    }
    
    public static class RSAUtil {
        private static final String ALGORITHM = "RSA";
        private static final int KEY_SIZE = 2048;
        
        /**
         * 生成RSA密钥对
         */
        public static KeyPair generateKeyPair() throws NoSuchAlgorithmException {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(ALGORITHM);
            keyPairGenerator.initialize(KEY_SIZE);
            return keyPairGenerator.generateKeyPair();
        }
        
        /**
         * 生成Base64编码的公私钥字符串
         * @return List中index 0为公钥，1为私钥
         */
        public static List<String> generateKeyStrings() throws NoSuchAlgorithmException {
            KeyPair keyPair = generateKeyPair();
            List<String> keyStrings = new ArrayList<>();
            
            String publicKey = Base64.getEncoder().encodeToString(keyPair.getPublic().getEncoded());
            String privateKey = Base64.getEncoder().encodeToString(keyPair.getPrivate().getEncoded());
            
            keyStrings.add(publicKey);
            keyStrings.add(privateKey);
            return keyStrings;
        }
        
        /**
         * 处理密钥字符串，移除PEM格式的头尾
         */
        private static String formatKey(String key) {
            return key.replaceAll("-----BEGIN (.*)-----", "")
                    .replaceAll("-----END (.*)-----", "")
                    .replaceAll("\\s+", "");
        }
        
        /**
         * 使用公钥加密
         * @param publicKeyStr Base64编码的公钥
         * @param plainText 明文
         * @return Base64编码的密文
         */
        public static String encrypt(String publicKeyStr, String plainText) throws Exception {
            publicKeyStr = formatKey(publicKeyStr);
            byte[] keyBytes = Base64.getDecoder().decode(publicKeyStr);
            
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
            PublicKey publicKey = keyFactory.generatePublic(keySpec);
            
            Cipher cipher = Cipher.getInstance("RSA/ECB/OAEPWITHSHA-1ANDMGF1PADDING");
            OAEPParameterSpec oaepParams = new OAEPParameterSpec("SHA-1", "MGF1",
                    MGF1ParameterSpec.SHA1, PSource.PSpecified.DEFAULT);
            cipher.init(Cipher.ENCRYPT_MODE, publicKey, oaepParams);
            
            byte[] encryptedBytes = cipher.doFinal(plainText.getBytes());
            return Base64.getEncoder().encodeToString(encryptedBytes);
        }
        
        /**
         * 使用私钥解密
         * @param privateKeyStr Base64编码的私钥
         * @param encryptedText Base64编码的密文
         * @return 明文
         */
        public static String decrypt(String privateKeyStr, String encryptedText) throws Exception {
            privateKeyStr = formatKey(privateKeyStr);
            byte[] keyBytes = Base64.getDecoder().decode(privateKeyStr);
            
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
            PrivateKey privateKey = keyFactory.generatePrivate(keySpec);
            
            Cipher cipher = Cipher.getInstance("RSA/ECB/OAEPWITHSHA-1ANDMGF1PADDING");
            OAEPParameterSpec oaepParams = new OAEPParameterSpec("SHA-1", "MGF1",
                    MGF1ParameterSpec.SHA1, PSource.PSpecified.DEFAULT);
            cipher.init(Cipher.DECRYPT_MODE, privateKey, oaepParams);
            
            byte[] decryptedBytes = cipher.doFinal(Base64.getDecoder().decode(encryptedText));
            return new String(decryptedBytes);
        }
    }
    
}
